import cv2
import numpy as np
import matplotlib.pyplot as plt
import math
import random


def Correlate(rawImg, mask):
	maskRadius = 0 
	maskShiftRange = 0


	maskRow, maskCol = mask.shape
	if (maskRow%2 == 0):
		maskRadius = maskRow / 2
		maskShiftRange = maskRadius
	else:
		maskRadius = (maskRow - 1) / 2
		maskShiftRange = maskRadius + 1


	imgHeight, imgWidth = rawImg.shape
	resultImage = rawImg.copy()

	#i = 0
	#j = 2
	for i in range(imgHeight):
		for j in range(imgWidth):
	#if True:
	#	if True:
			tempSum = 0.0
			for xDelta in range(-maskRadius, maskShiftRange):
				for yDelta in range(-maskRadius, maskShiftRange):
					imgX = i + xDelta
					imgY = j + yDelta
					maskX = xDelta + maskRadius
					maskY = yDelta + maskRadius
					
					if ((imgX < 0) or (imgY < 0) or (imgX >= imgHeight) or (imgY >= imgWidth)):
						continue
					else:
						#print rawImg[imgX, imgY], mask[maskX, maskY]
						tempSum += rawImg[imgX, imgY] * mask[maskX, maskY]

					#print tempSum
			resultImage[i,j] = abs(int(round(tempSum)))

	return resultImage


def ApplyLaplacianMask(rawImg, useEnhancedMask = False, threshold = 122, displayIntermediaResult = False):
	mask = np.zeros([3, 3])
	if useEnhancedMask:
		mask[0, :] = [1, 1, 1]
		mask[1, :] = [1, -8, 1]
		mask[2, :] = [1, 1, 1]
	else:
		mask[0, :] = [0, 1, 0]
		mask[1, :] = [1, -4, 1]
		mask[2, :] = [0, 1, 0]

	
	imgHeight, imgWidth = rawImg.shape

	edge = Correlate(rawImg, mask)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (edge[i, j] <= threshold):
				edge[i, j] = 0	

	return edge


def ApplyRobertsMask(rawImg, threshold = 122, displayIntermediaResult = False):
	maskA = np.zeros([2, 2])
	maskA[0, :] = [-1, 0]
	maskA[1, :] = [0, 1]

	maskB = np.zeros([2, 2])
	maskB[0, :] = [0, -1]
	maskB[1, :] = [1, 0]
	imgHeight, imgWidth = rawImg.shape

	horizontalEdge = Correlate(rawImg, maskA)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (horizontalEdge[i, j] <= threshold):
				horizontalEdge[i, j] = 0

	if displayIntermediaResult:
		#print horizontalEdge
		cv2.namedWindow('Rebert \ Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Rebert \ Result", horizontalEdge)
	
	verticalEdge = Correlate(rawImg, maskB)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (verticalEdge[i, j] <= threshold):
				verticalEdge[i, j] = 0


	if displayIntermediaResult:
		#print verticalEdge
		cv2.namedWindow('Rebert / Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Rebert / Result", verticalEdge)

	result = verticalEdge + horizontalEdge

	return result


def ApplyPrewittMask(rawImg, threshold = 122, displayIntermediaResult = False):
	maskH = np.zeros([3, 3])
	maskH[0, :] = [1, 1, 1]
	maskH[2, :] = [-1, -1, -1]

	maskV = np.zeros([3, 3])
	maskV[0, :] = [-1, 0, 1]
	maskV[1, :] = [-1, 0, 1]
	maskV[2, :] = [-1, 0, 1]

	imgHeight, imgWidth = rawImg.shape

	horizontalEdge = Correlate(rawImg, maskH)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (horizontalEdge[i, j] <= threshold):
				horizontalEdge[i, j] = 0

	if displayIntermediaResult:
		#print horizontalEdge
		cv2.namedWindow('Prewitt Horizontal Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Prewitt Horizontal Result", horizontalEdge)
	
	verticalEdge = Correlate(rawImg, maskV)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (verticalEdge[i, j] <= threshold):
				verticalEdge[i, j] = 0


	if displayIntermediaResult:
		#print verticalEdge
		cv2.namedWindow('Prewitt Vertical Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Prewitt Vertical Result", verticalEdge)

	result = verticalEdge + horizontalEdge

	return result

def ApplySobelMask(rawImg, threshold = 122, displayIntermediaResult = False):
	maskV = np.zeros([3, 3])
	maskV[0, :] = [-1, -2, -1]
	maskV[2, :] = [1, 2, 1]
	
	maskH = np.zeros([3, 3])
	maskH[0, :] = [-1, 0, 1]
	maskH[1, :] = [-2, 0, 2]
	maskH[2, :] = [-1, 0, 1]

	imgHeight, imgWidth = rawImg.shape

	horizontalEdge = Correlate(rawImg, maskH)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (horizontalEdge[i, j] <= threshold):
				horizontalEdge[i, j] = 0

	if displayIntermediaResult:
		#print horizontalEdge
		cv2.namedWindow('Sobel Horizontal Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Sobel Horizontal Result", horizontalEdge)
	
	verticalEdge = Correlate(rawImg, maskV)

	for i in range(imgHeight):
		for j in range(imgWidth):
			if (verticalEdge[i, j] <= threshold):
				verticalEdge[i, j] = 0


	if displayIntermediaResult:
		#print verticalEdge
		cv2.namedWindow('Sobel Vertical Result', cv2.CV_WINDOW_AUTOSIZE)
		cv2.imshow( "Sobel Vertical Result", verticalEdge)

	result = verticalEdge + horizontalEdge

	return result


if __name__ == '__main__':
	print "========== Test ========== "
	# testData = np.ones((5, 5))
	# testResult = ApplySobelMask(testData, 0)
	# print testResult

	testData = np.zeros((5, 5))
	testData[2, :] = 1
	testData[:, 2] = 1
	print testData
	testResult = ApplySobelMask(testData, 0, False)
	print testResult

	testData = np.ones((5, 5))
	testData[2, :] = 0
	testData[:, 2] = 0
	print testData
	testResult = ApplySobelMask(testData, 0)
	print testResult

	print "========== Loading Test Image ========== "
	rawImg = cv2.imread('./img/fig3.tif',0)
	# rawImg = cv2.imread('./img/Bikesgray.jpg',0)
	cv2.namedWindow('Raw Image', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Raw Image", rawImg)

	print "========== Processing Fig3 by Laplacian ========== "
	lapResult = ApplyLaplacianMask(rawImg)
	cv2.namedWindow('Laplacian Result', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Laplacian Result", lapResult)

	print "========== Processing Fig3 by Enhanced Laplacian ========== "
	enLapResult = ApplyLaplacianMask(rawImg, True)
	cv2.namedWindow('Enhanced Laplacian Result', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Enhanced Laplacian Result", enLapResult)

	print "========== Processing Fig3 by Robert ========== "
	robertResult = ApplyRobertsMask(rawImg, 0) #, True)
	cv2.namedWindow('Robert Result', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Robert Result", robertResult)

	print "========== Processing Fig3 by Prewitt ========== "
	PrewittResult = ApplyPrewittMask(rawImg) #122, True)
	cv2.namedWindow('Prewitt Result', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Prewitt Result", PrewittResult)

	print "========== Processing Fig3 by Sobel ========== "
	sobelResult = ApplySobelMask(rawImg)
	cv2.namedWindow('Sobel Result', cv2.CV_WINDOW_AUTOSIZE)
	cv2.imshow( "Sobel Result", sobelResult)

	print "========== Done ========== "


	cv2.waitKey(0)
	cv2.destroyAllWindows()
